# 此函数用于对称分布傅里叶的色散部分的仿真，GVD+PMD
# 传入参数：
#   polar_n:偏振数
#   beta_2:群速度色散,s^2/km
#   f_omega
#   deltaz:每个span进行分布傅里叶的长度
#   PMD:是否添加偏振模色散
#   PMD_str_para:PDM构造好的参数

import numpy as np
from numpy.fft import fft, ifft ,fftshift, ifftshift


def Dispersion(signal, Signal, Fiber_para, PMD_para, Dis):
    data = signal
    '=================================================信号相关参数'
    npol = Signal['polarization']
    lamda_wave = Signal['lamda_wave']
    field = Signal['field']
    C = Signal['C']
    f_omega = Signal['f_omega']
    # from scipy import io
    # f_omega = io.loadmat('f_omgea.mat')
    # f_omega = f_omega['f_omega']
    # 双偏振情况
    if npol == 2:
        # x偏振信号
        data_x = data[0:, 0]
        data_fft_x = fftshift(fft(data_x))
        # y偏振信号
        data_y = data[0:, 1]
        data_fft_y = fftshift(fft(data_y))

        data_fft_x = data_fft_x.reshape(-1, 1)
        data_fft_y = data_fft_y.reshape(-1, 1)
        data_fft = np.concatenate((data_fft_x, data_fft_y), axis=1)
        # data_fft = data_fft.reshape((data_fft[0].shape[0], len(data_fft)))
    else:
        data = fft(data)
    '================================================光纤相关参数'
    #  alpha_db = Fiber_para.alpha_dB
    #  db/km--->1/km
    #  alpha_km = Fiber_para.alpha_km
    L = Fiber_para['L']
    lamda_Fiber = Fiber_para['lambda']
    StepSize = Fiber_para['step_size']
    '================================================偏振模相关参数'
    PMD = PMD_para['PMD']
    #    PMD_DGD = PMD_para.PMD_DGD
    #    PMD_DGD_Trunk = PMD_para.PMD_DGD_Trunk
    #    Trunk_Num = PMD_para.trunk_num
    #    trunk_length = PMD_para.trunk_length
    # Trunk_Num = PMD_para['trunk_num']
    L_trunk = PMD_para['L_trunk']
    Index_trunk = PMD_para['Index_trunk']
    Trunk_Matrix = PMD_para['Matrix']
    # 特征值deltabeta0 [1 / m]
    eigenvalue = PMD_para['eigenvalue']
    # [1/km]
    eigenvalue = eigenvalue * 1e3

    '================================================差分群时延相关参数'
    # [ns/km]
    DGD = PMD_para['DGD_unit']
    # [s/km]
    DGD = DGD * 1e-9

    # 群速度及高阶色散参数
    beta_2_fiber = Fiber_para['beta_2']  # [s^2/km]
    beta_3_fiber = Fiber_para['beta_3']  # [s^3/km]

    # 总色散
    if npol == 2:
        # beta = np.zeros((len(data_x), npol))
        beta = np.zeros((len(data_x), npol))
    else:
        beta = np.zeros((len(data), npol))

    # 主体部分
    # dual-polarization
    if npol == 2:
        # add PMD
        if PMD == 1:
            # PMD
            # GVD and TOD
            beta_2 = beta_2_fiber  # beta2 [s^2/km]
            beta_3 = beta_3_fiber  # [s^3/km]
            beta_ = f_omega * (f_omega * (beta_2 / 2 + f_omega * beta_3 / 6))
            beta_ = beta_.reshape(-1, 1)
            # beta_ = f_omega * (beta_1 + (f_omega * (beta_2 / 2)))
            # x和y的色散[1 / km]
            beta[0:, 0] = beta_[0:,0]
            f_omega_ = f_omega.reshape(-1, 1)
            beta[0:, 1] = beta_[0:,0] - DGD * f_omega_[0:,0]
            for i in range(len(L_trunk)):
                # 随机旋转矩阵
                Matrix = Trunk_Matrix[:, :, int(Index_trunk[i] - 1)]

                Matrix_H = np.transpose(Trunk_Matrix[:, :, int(Index_trunk[i] - 1)])

                data_fft = data_fft @ np.conj(Matrix)
                data_fft = data_fft * np.exp(-1j * beta * StepSize / 2)
                if Dis == 2:
                    eigenvalue_ = eigenvalue[int(Index_trunk[i] - 1), :]
                    data_fft[0:, 0] = data_fft[0:, 0] * np.exp(-1j * eigenvalue_[0] * L_trunk[i])
                    data_fft[0:, 1] = data_fft[0:, 1] * np.exp(-1j * eigenvalue_[1] * L_trunk[i])
                data_fft = data_fft @ Matrix_H

            sig_x = ifft(ifftshift(data_fft[0:, 0]))
            sig_y = ifft(ifftshift(data_fft[0:, 1]))

            sig_x = sig_x.reshape(-1, 1)
            sig_y = sig_y.reshape(-1, 1)

            sig = np.concatenate((sig_x, sig_y), axis=1)

        else:
            # NO PMD
            # GVD and TOD
            # print('色散前功率')
            # print(np.mean(np.abs(data_fft[0:, 0]) ** 2))
            beta_2 = beta_2_fiber  # beta2 [s^2/km]
            # beta_2 = -21.6676e-24
            beta_3 = beta_3_fiber  # [s^3/km]
            beta_ = f_omega * (f_omega * (beta_2 / 2 + f_omega * beta_3 / 6))
            # beta_ = f_omega * (beta_1 + (f_omega * (beta_2 / 2)))
            # x和y的色散[1 / km]
            beta_ = beta_.reshape(-1, 1)
            beta[0:, 0] = beta_[0:,0]
            beta[0:, 1] = beta_[0:,0]
            rows, cols = data_fft.shape
            beta = beta.reshape(rows, cols)

            data_fft[0:, 0] = data_fft[0:, 0] * np.exp(-1j * beta[0:, 0] * StepSize / 2)
            data_fft[0:, 1] = data_fft[0:, 1] * np.exp(-1j * beta[0:, 1] * StepSize / 2)


            sig_x = ifft(ifftshift(data_fft[0:, 0]))
            sig_y = ifft(ifftshift(data_fft[0:, 1]))

            sig_x = sig_x.reshape(-1, 1)
            sig_y = sig_y.reshape(-1, 1)

            sig = np.concatenate((sig_x, sig_y), axis=1)
            # print('色散后功率')
            # print(np.mean(np.abs(sig_x) ** 2))
    else:
        # single-polarization
        beta_2 = beta_2_fiber  # beta2 [s^2/km]
        beta_3 = beta_3_fiber  # [s^3/km]
        beta_ = f_omega * (f_omega * (beta_2 / 2 + f_omega * beta_3 / 6))
        # beta_ = f_omega * (beta_1 + (f_omega * (beta_2 / 2)))
        # 色散[1 / km]
        data_fft = data_fft * np.exp(-1j * beta_ * StepSize / 2)
        sig = ifft(ifftshift(data_fft))

    return sig
